// Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//  * Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//  * Neither the name of NVIDIA CORPORATION nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "imageembedder.h"

#include <getopt.h>
#include <unistd.h>
#include <cassert>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <sstream>

#include <opencv4/opencv2/opencv.hpp>


namespace tc = triton::client;

#define FAIL_IF_ERR(X, MSG)                                        \
  {                                                                \
    tc::Error err = (X);                                           \
    if (!err.IsOk()) {                                             \
      std::cerr << "error: " << (MSG) << ": " << err << std::endl; \
      exit(1);                                                     \
    }                                                              \
  }

namespace {

void
ValidateShapeAndDatatype(
    const std::string& name, std::shared_ptr<tc::InferResult> result)
{
  std::vector<int64_t> shape;
  FAIL_IF_ERR(
      result->Shape(name, &shape), "unable to get shape for '" + name + "'");
  // Validate shape
  if ((shape.size() != 1) || shape[0] != 512) {
    std::cerr << "error: received incorrect shapes for '" << name << "'"
              << std::endl;
    exit(1);
  }
  std::string datatype;
  FAIL_IF_ERR(
      result->Datatype(name, &datatype),
      "unable to get datatype for '" + name + "'");
  // Validate datatype
  if (datatype.compare("FP32") != 0) {
    std::cerr << "error: received incorrect datatype for '" << name
              << "': " << datatype << std::endl;
    exit(1);
  }
}

}  // namespace

const std::string base64_chars =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz"
    "0123456789+/";

std::string
base64Encode(const unsigned char* bytes_to_encode, unsigned int in_len)
{
  std::string ret;
  int i = 0;
  int j = 0;
  unsigned char char_array_3[3];
  unsigned char char_array_4[4];

  while (in_len--) {
    char_array_3[i++] = *(bytes_to_encode++);
    if (i == 3) {
      char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
      char_array_4[1] =
          ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
      char_array_4[2] =
          ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
      char_array_4[3] = char_array_3[2] & 0x3f;
      for (i = 0; (i < 4); i++) {
        ret += base64_chars[char_array_4[i]];
      }
      i = 0;
    }
  }

  if (i) {
    for (j = i; j < 3; j++) {
      char_array_3[j] = '\0';
    }
    char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
    char_array_4[1] =
        ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
    char_array_4[2] =
        ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
    char_array_4[3] = char_array_3[2] & 0x3f;
    for (j = 0; (j < i + 1); j++) {
      ret += base64_chars[char_array_4[j]];
    }
    while ((i++ < 3)) {
      ret += '=';
    }
  }

  return ret;
}

std::vector<float>
imageToFloatVector(const cv::Mat& image)
{
  std::vector<float> floatVector;

  // 确保图像不为空且为三通道图像
  if (image.empty() || image.channels() != 3) {
    std::cerr << "Image is empty or not 3 channels, "
                 "can't convert to vector"
              << std::endl;
    return floatVector;
  }

  // 将每个通道的像素值添加到向量中
  for (int c = 0; c < image.channels(); ++c) {
    for (int i = 0; i < image.rows; ++i) {
      for (int j = 0; j < image.cols; ++j) {
        float pixelValue = image.at<cv::Vec3f>(i, j)[c];
        floatVector.push_back(pixelValue);
      }
    }
  }

  return floatVector;
}

// 去除 Base64 字符串中的前缀 "data:image/png;base64,"
std::string
cleanBase64String(const std::string& base64Str)
{
  return base64Str.substr(base64Str.find(",") + 1);
}

std::vector<unsigned char>
base64Decode(const std::string& input)
{
  std::vector<unsigned char> output;
  const std::string base64Chars =
      "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
      "abcdefghijklmnopqrstuvwxyz"
      "0123456789+/";

  int val = 0, valb = -8;
  for (auto c : input) {
    if (std::isspace(c)) {
      continue;
    }
    if (c == '=') {
      break;
    }
    const auto ch = std::find(base64Chars.begin(), base64Chars.end(), c);
    if (ch == base64Chars.end()) {
      throw std::runtime_error("Invalid base64 character");
    }

    val = (val << 6) + (ch - base64Chars.begin());
    valb += 6;
    if (valb >= 0) {
      output.push_back(static_cast<unsigned char>((val >> valb) & 0xFF));
      valb -= 8;
    }
  }

  return output;
}

namespace math {

// 向量取模
float
norm(const std::vector<float>& vec)
{
  float norm = 0.0f;
  for (const float& element : vec) {
    norm += element * element;
  }
  norm = std::sqrt(norm);

  return norm;
}

void
normalize(std::vector<float>& vec)
{
  if (vec.empty()) {
    return;
  }

  float normNum = norm(vec);
  if (normNum == 0.0) {
    return;
  }

  for (int i = 0; i < vec.size(); ++i) {
    vec[i] = vec[i] / normNum;
  }
}
}  // namespace math

ImageEmbedder::ImageEmbedder()
{
  long verify_peer = 1;
  long verify_host = 2;
  std::string cacerts;
  std::string certfile;
  std::string keyfile;
  tc::HttpSslOptions sslOptions;
  sslOptions.verify_peer = verify_peer;
  sslOptions.verify_host = verify_host;
  sslOptions.ca_info = cacerts;
  sslOptions.cert = certfile;
  sslOptions.key = keyfile;

  FAIL_IF_ERR(tc::InferenceServerHttpClient::Create(&httpClient_, httpServerUrl_, true /* verbose */, sslOptions),
  "unable to create http client");
}

std::vector<float> ImageEmbedder::embedImage(const std::string &imagePath)
{
  tc::Headers http_headers;
  uint32_t clientTimeout = 0;
  auto request_compression_algorithm =
      tc::InferenceServerHttpClient::CompressionType::NONE;
  auto response_compression_algorithm =
      tc::InferenceServerHttpClient::CompressionType::NONE;

  std::vector<unsigned char> buffer = readImageBuffer(imagePath);

  // 这里模拟了 base64 编码的传输，以前 ai-business 里接收的是 base64 编码格式的图像，应用需要传给它这种格式。
  // 这里模拟了一下，先把图像 encode 成 base64，再 decode 回图像。
  std::string base64EncodedString = base64Encode(buffer.data(), buffer.size());
  std::vector<unsigned char> image = base64Decode(base64EncodedString);

  std::vector<float> input0_data = preprocessVision(image);

  std::vector<int64_t> shape = {3, 224, 224};

  // Initialize the inputs with the data.
  tc::InferInput *input0;

  FAIL_IF_ERR(
      tc::InferInput::Create(&input0, "image", shape, "FP32"),
      "unable to get image");
  std::shared_ptr<tc::InferInput> input0_ptr;
  input0_ptr.reset(input0);

  FAIL_IF_ERR(
      input0_ptr->AppendRaw(
          reinterpret_cast<uint8_t *>(&input0_data[0]),
          input0_data.size() * sizeof(float)),
      "unable to set data for image");

  // The inference settings. Will be using default for now.
  tc::InferOptions options(modelName_);
  options.model_version_ = modelVersion_;
  options.client_timeout_ = clientTimeout;
  options.priority_ = 1;

  std::vector<tc::InferInput *> inputs = {input0_ptr.get()};
  // Empty output vector will request data for all the output tensors from
  // the server.
  std::vector<const tc::InferRequestedOutput *> outputs = {};

  tc::InferResult *results;
  FAIL_IF_ERR(
      httpClient_->Infer(
          &results, options, inputs, outputs, http_headers, tc::Parameters(),
          request_compression_algorithm, response_compression_algorithm),
      "unable to run model");
  std::shared_ptr<tc::InferResult> results_ptr;
  results_ptr.reset(results);

  // Validate the results...
  ValidateShapeAndDatatype("unnorm_image_features", results_ptr);

  // Get pointers to the result returned...
  float *output0_data;
  size_t output0_byte_size;
  FAIL_IF_ERR(
      results_ptr->RawData(
          "unnorm_image_features", (const uint8_t **)&output0_data,
          &output0_byte_size),
      "unable to get result data for 'unnorm_image_features'");
  if (output0_byte_size != 512 * sizeof(float))
  {
    std::cerr
        << "error: received incorrect byte size for 'unnorm_image_features': "
        << output0_byte_size << std::endl;
    exit(1);
  }

  std::vector<float> output0_data_vector(output0_data, output0_data + 512);
  math::normalize(output0_data_vector);

  // Get full response
  std::cout << results_ptr->DebugString() << std::endl;

  tc::InferStat infer_stat;
  httpClient_->ClientInferStat(&infer_stat);
  std::cout << "======Client Statistics======" << std::endl;
  std::cout << "completed_request_count " << infer_stat.completed_request_count
            << std::endl;
  std::cout << "cumulative_total_request_time_ns "
            << infer_stat.cumulative_total_request_time_ns << std::endl;
  std::cout << "cumulative_send_time_ns " << infer_stat.cumulative_send_time_ns
            << std::endl;
  std::cout << "cumulative_receive_time_ns "
            << infer_stat.cumulative_receive_time_ns << std::endl;

  std::string model_stat;
  FAIL_IF_ERR(
      httpClient_->ModelInferenceStatistics(&model_stat, modelName_),
      "unable to get model statistics");
  std::cout << "======Model Statistics======" << std::endl;
  std::cout << model_stat << std::endl;

  std::cout << "PASS : Infer" << std::endl;

  return output0_data_vector;
}

std::vector<unsigned char> ImageEmbedder::readImageBuffer(const std::string &imagePath)
{
    // 读取文件内容
  std::ifstream fileStream(imagePath, std::ios::binary);
  if (!fileStream.is_open()) {
    std::cerr << "Failed to open file: " << imagePath << std::endl;
    return std::vector<unsigned char>();
  }

  // 获取文件大小
  fileStream.seekg(0, std::ios::end);
  std::streampos fileSize = fileStream.tellg();
  fileStream.seekg(0, std::ios::beg);

  // 读取文件内容到 vector
  std::vector<unsigned char> buffer(fileSize);
  fileStream.read(reinterpret_cast<char*>(buffer.data()), fileSize);
  fileStream.close();

  return buffer;
}


std::vector<float>
ImageEmbedder::preprocessVision(const std::vector<unsigned char>& base64Image)
{
  const int width = 224;
  const int height = 224;
  const cv::Scalar mean = cv::Scalar(0.48145466, 0.4578275, 0.40821073);
  const cv::Scalar deviation = cv::Scalar(0.26862954, 0.26130258, 0.27577711);

  cv::Mat image = cv::imdecode(base64Image, cv::IMREAD_COLOR);
  if (image.empty()) {
    std::cerr << "Can't read image " << std::endl;
    return std::vector<float>();
  }

  cv::Size inputSize(width, height);
  cv::resize(image, image, inputSize, 0, 0, cv::INTER_CUBIC);
  cv::cvtColor(image, image, cv::COLOR_BGR2RGB);

  // Normalize
  image.convertTo(image, CV_32FC3);
  image /= 255.0;
  image -= mean;
  image /= deviation;

  // Convert to PyTorch tensor format by transposing the channels
  cv::transpose(image, image);

  return imageToFloatVector(image);
}